package fem2.strategies;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;
import fem2.Model;
import fem2.enu.MassMatrixType;
import fem2.jlinalg.DiagonalMatrix;
import fem2.jlinalg.Solver;

public abstract class CDMTimeIntegrationStrategy extends ExplicitTimeIntegrationStrategy {

	private double stepping;

	public CDMTimeIntegrationStrategy(Model m, MassMatrixType massType, double alpha1,
			double alpha2, double t0, double stepping, int nt) {
		super(m, massType, alpha1, alpha2);
		this.t = t0;
		this.stepping = stepping;
		this.setNumberOfTimeStep(nt);
	}

	@Override
	public void initializeSystemMatricesAndVectors(Solver leftHandSideSolver,
			ArrayVector rightHandSide, ArrayVector u) {
		int n = model.getSize();

		initializeSolver(leftHandSideSolver);

		rightHandSide.setSize(n);

		u.setSize(n);
	}

	@Override
	public void started(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		super.started(leftHandSideSolver, rightHandSide, u);

		/*
		 * compute uHat_old
		 */
		uHat_old = (ArrayVector) uHat.clone();
		BLAM.add(-stepping, udHat, uHat_old);
		BLAM.add(Math.pow(stepping, 2) / 2, uddHat, uHat_old);

		udHat_old = (ArrayVector) udHat.clone();
		uddHat_old = (ArrayVector) uddHat.clone();
	}

	@Override
	public void assembleEffectiveSystem(IMatrix kteff, IVector reff) {
		int n = model.getSize();
		double dt = t - t_old;

		/*
		 * assemble kt and M
		 */
		BLAM.zero(kteff);

		IMatrix M;
		if (massType == MassMatrixType.DIRECT_MASS_LUMPING) {
			M = new DiagonalMatrix(n);
		} else {
			M = (IMatrix) kteff.clone();
		}

		model.assembleKt(kteff);
		model.assembleM(M, massType);

		/*
		 * assemble reff
		 */
		model.assembleR(1, ((ArrayVector) reff).getValues(), t);

		ArrayVector ri = new ArrayVector(n);
		model.assembleRi(ri.getValues());
		BLAM.add(-1.0, ri, reff);

		BLAM.multiply(2 / Math.pow(dt, 2), BLAM.NO_TRANSPOSE, M, uHat, 1.0, reff);

		BLAM.multiply(alpha1 / (2 * dt) - 1 / Math.pow(dt, 2), BLAM.NO_TRANSPOSE, M, uHat_old, 1.0,
				reff);

		BLAM.multiply(alpha2 / (2 * dt), BLAM.NO_TRANSPOSE, kteff, uHat_old, 1.0, reff);

		/*
		 * assemble kteff
		 */
		BLAM.scale(alpha2 / (2 * dt), kteff);
		BLAM.add(1 / Math.pow(dt, 2) + alpha1 / (2 * dt), M, kteff);

	}

	@Override
	public void updateSolution(IVector U) {
		/*
		 * update udHat_old. For this scheme, velocity and acceleration are one
		 * step behind the displacement
		 */
		BLAM.copy(U, udHat_old);
		BLAM.add(-1.0, uHat_old, udHat_old);
		BLAM.scale(1 / (2 * stepping), udHat_old);
		/*
		 * update uddHat_old
		 */
		BLAM.copy(U, uddHat_old);
		BLAM.add(-2.0, uHat, uddHat_old);
		BLAM.add(1.0, uHat_old, uddHat_old);
		BLAM.scale(1 / Math.pow(stepping, 2), uddHat_old);
		/*
		 * update un_minus_1
		 */
		BLAM.copy(uHat, uHat_old);
		/*
		 * update un
		 */
		BLAM.copy(U, uHat);
	}

	@Override
	public void timeStepFinished(Solver leftHandSideSolver, ArrayVector rightHandSide, ArrayVector u) {
		/*
		 * trigger model to act on material
		 */
		model.FinalizeSolutionStep();
	}

	@Override
	public double getNextTimeStep() {
		t += stepping;
		return t;
	}

	@Override
	public double getCurrentTimeStep() {
		return t;
	}

}
